"""
本地视觉模型集成模块
整合LocalVisionModel和VisionModelOptimizer，提供统一的本地视觉处理接口
"""

import logging
import asyncio
from typing import Optional, Dict, Any, Union
from concurrent.futures import Future
import numpy as np

from .local_vision_model import LocalVisionModel, create_local_vision_model
from .vision_model_optimizer import VisionModelOptimizer, create_vision_optimizer


class OptimizedLocalVisionModel:
    """
    优化的本地视觉模型
    结合LocalVisionModel和VisionModelOptimizer的功能
    """
    
    def __init__(self, model_path: str, model_type: Optional[str] = None, 
                 optimization_config: Optional[Dict[str, Any]] = None):
        """
        初始化优化的本地视觉模型
        
        Args:
            model_path: 模型文件路径
            model_type: 模型类型
            optimization_config: 优化配置
        """
        self.logger = logging.getLogger(__name__)
        
        # 创建基础模型
        self.base_model = create_local_vision_model(model_path, model_type)
        
        # 创建优化器
        self.optimizer = create_vision_optimizer(optimization_config)
        
        # 优化模型
        if hasattr(self.base_model.backend, 'model') and self.base_model.backend.model is not None:
            self.base_model.backend.model = self.optimizer.optimize_model(
                self.base_model.backend.model, 
                self.base_model.model_type
            )
        
        self.logger.info("Optimized local vision model initialized")
    
    def analyze_image(self, image: np.ndarray, query: str = "Describe this image") -> str:
        """
        同步图像分析
        
        Args:
            image: 输入图像
            query: 查询文本
            
        Returns:
            分析结果
        """
        return self.base_model.analyze_image(image, query)
    
    async def analyze_image_async(self, image: np.ndarray, query: str = "Describe this image", 
                                 priority: int = 0) -> str:
        """
        异步图像分析（支持批处理优化）
        
        Args:
            image: 输入图像
            query: 查询文本
            priority: 请求优先级
            
        Returns:
            分析结果
        """
        try:
            # 提交到优化器进行批处理
            future = self.optimizer.submit_inference_request(image, query, priority)
            
            # 等待结果
            result = await asyncio.wrap_future(future)
            
            if result.success:
                return result.result
            else:
                # 如果批处理失败，回退到基础模型
                self.logger.warning(f"Batch processing failed: {result.error_message}")
                return self.base_model.analyze_image(image, query)
                
        except Exception as e:
            self.logger.error(f"Async analysis failed: {e}")
            # 回退到同步处理
            return self.base_model.analyze_image(image, query)
    
    def get_model_info(self) -> Dict[str, Any]:
        """获取模型信息"""
        base_info = self.base_model.get_model_info()
        optimization_info = self.optimizer.get_optimization_info()
        
        return {
            'base_model': base_info,
            'optimization': optimization_info,
            'features': {
                'async_processing': True,
                'batch_processing': optimization_info['config'].get('enable_batch_processing', False),
                'gpu_acceleration': optimization_info['config'].get('enable_gpu_acceleration', False),
                'quantization': optimization_info['config'].get('enable_quantization', False)
            }
        }
    
    def is_model_loaded(self) -> bool:
        """检查模型是否已加载"""
        return self.base_model.is_model_loaded()
    
    def shutdown(self):
        """关闭模型和优化器"""
        self.optimizer.shutdown()
        self.logger.info("Optimized local vision model shutdown")


class LocalVisionService:
    """
    本地视觉服务
    提供高级的本地视觉处理服务接口
    """
    
    def __init__(self, models_config: Dict[str, Dict[str, Any]]):
        """
        初始化本地视觉服务
        
        Args:
            models_config: 模型配置字典
                格式: {
                    'model_name': {
                        'path': 'model_path',
                        'type': 'model_type',
                        'optimization': {...}
                    }
                }
        """
        self.logger = logging.getLogger(__name__)
        self.models = {}
        self.default_model = None
        
        # 加载配置的模型
        for model_name, config in models_config.items():
            try:
                model = OptimizedLocalVisionModel(
                    model_path=config['path'],
                    model_type=config.get('type'),
                    optimization_config=config.get('optimization')
                )
                
                self.models[model_name] = model
                
                # 设置第一个成功加载的模型为默认模型
                if self.default_model is None:
                    self.default_model = model_name
                    
                self.logger.info(f"Loaded model: {model_name}")
                
            except Exception as e:
                self.logger.error(f"Failed to load model {model_name}: {e}")
        
        if not self.models:
            raise RuntimeError("No models could be loaded")
    
    def analyze_image(self, image: np.ndarray, query: str = "Describe this image", 
                     model_name: Optional[str] = None) -> str:
        """
        使用指定模型分析图像
        
        Args:
            image: 输入图像
            query: 查询文本
            model_name: 模型名称，如果为None则使用默认模型
            
        Returns:
            分析结果
        """
        model_name = model_name or self.default_model
        
        if model_name not in self.models:
            raise ValueError(f"Model {model_name} not found")
        
        return self.models[model_name].analyze_image(image, query)
    
    async def analyze_image_async(self, image: np.ndarray, query: str = "Describe this image",
                                 model_name: Optional[str] = None, priority: int = 0) -> str:
        """
        异步图像分析
        
        Args:
            image: 输入图像
            query: 查询文本
            model_name: 模型名称
            priority: 请求优先级
            
        Returns:
            分析结果
        """
        model_name = model_name or self.default_model
        
        if model_name not in self.models:
            raise ValueError(f"Model {model_name} not found")
        
        return await self.models[model_name].analyze_image_async(image, query, priority)
    
    def get_available_models(self) -> Dict[str, Dict[str, Any]]:
        """获取可用模型列表"""
        return {
            name: model.get_model_info() 
            for name, model in self.models.items()
        }
    
    def set_default_model(self, model_name: str):
        """设置默认模型"""
        if model_name not in self.models:
            raise ValueError(f"Model {model_name} not found")
        
        self.default_model = model_name
        self.logger.info(f"Default model set to: {model_name}")
    
    def shutdown(self):
        """关闭所有模型"""
        for model in self.models.values():
            model.shutdown()
        self.logger.info("Local vision service shutdown")


def create_local_vision_service(models_config: Dict[str, Dict[str, Any]]) -> LocalVisionService:
    """
    工厂函数：创建本地视觉服务
    
    Args:
        models_config: 模型配置
        
    Returns:
        LocalVisionService实例
    """
    return LocalVisionService(models_config)


# 示例配置
EXAMPLE_CONFIG = {
    'default_onnx': {
        'path': '/path/to/model.onnx',
        'type': 'onnx',
        'optimization': {
            'enable_gpu_acceleration': True,
            'enable_quantization': True,
            'enable_batch_processing': True,
            'batch_size': 4
        }
    },
    'pytorch_model': {
        'path': '/path/to/pytorch/model',
        'type': 'pytorch',
        'optimization': {
            'enable_gpu_acceleration': True,
            'enable_quantization': True,
            'enable_batch_processing': True,
            'batch_size': 8
        }
    }
}


# 示例使用
if __name__ == "__main__":
    # 配置日志
    logging.basicConfig(level=logging.INFO)
    
    # 检查支持的格式
    from .local_vision_model import LocalVisionModel
    formats = LocalVisionModel.get_supported_formats()
    print("Supported formats:", formats)
    
    # 示例：创建服务（需要实际的模型文件）
    # service = create_local_vision_service(EXAMPLE_CONFIG)
    # result = service.analyze_image(image_array, "What do you see?")
    # print(result)
    # service.shutdown()